/*
 * Copyright (c) 2024 endless-sky
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include "upgrade.h"
#include "uboot_env_manager.h"
#include "ubi_manager.h"
#include "utils_noncopyable.h"
#include <memory>
#include <string>
#include <mutex>
#include <atomic>
#include <vector>
#include <map>
#include <thread>
#include <nlohmann/json.hpp>

namespace El {
namespace Upgrade {

// 前向声明
class MtdManager;

// CPIO格式相关结构体定义
struct CpioNewcHeader {
    char magic[6];     // "070701"
    char ino[8];       // inode number
    char mode[8];      // file mode
    char uid[8];       // user id
    char gid[8];       // group id
    char nlink[8];     // number of links
    char mtime[8];     // modification time
    char filesize[8];  // file size
    char devmajor[8];  // device major number
    char devminor[8];  // device minor number
    char rdevmajor[8]; // special device major number
    char rdevminor[8]; // special device minor number
    char namesize[8];  // filename length
    char check[8];     // checksum
    // filename follows (namesize bytes, null-terminated)
    // file data follows (filesize bytes)
};

struct CpioFileInfo {
    std::string filename;
    uint32_t filesize;
    uint32_t namesize;
    size_t headerOffset; // 头部在缓冲区中的偏移
    size_t dataOffset;   // 数据在缓冲区中的偏移
    size_t bytesRead;    // 已读取的文件数据字节数
};

/**
 * @brief 基于libmtd的流式升级实现
 *
 * 提供完整的流式升级功能：
 * - A/B分区无缝升级
 * - libmtd直接操作MTD设备
 * - 流式数据处理，边收边升级
 * - U-Boot环境变量管理
 * - 升级状态监控
 * - 错误处理和恢复
 */
class UpgradeImpl : public IUpgradeService {
    DISALLOW_COPY_AND_MOVE(UpgradeImpl);

public:
    static UpgradeImpl &GetInstance();

    // 实现 IUpgradeService 接口
    bool Start() override;
    bool InitializeUpgradeEnvironment() override;
    void AbortUpgrade() override;
    void RegisterProgressCallback(UpgradeProgressCallback callback) override;
    bool StreamUpgradeData(const uint8_t *data, size_t length) override;
    bool CompleteUpgrade() override;
    std::string GetBootPartitionInfo() override;
    bool SwitchToBootPartition(const std::string &slot) override;
    std::string GetCurrentUpgradeStatus() override;
    std::string GetLastErrorMessage() override;

private:
    UpgradeImpl();
    ~UpgradeImpl();

    /**
     * @brief 获取最优缓冲区大小
     * 基于可用内存动态调整缓冲区大小(1MB-4MB)
     * @return 最优缓冲区大小
     */
    size_t GetOptimalBufferSize() const;

    /**
     * @brief 处理升级包数据流
     * @param data 数据指针
     * @param length 数据长度
     * @return 是否成功
     */
    bool ProcessUpgradePackageData(const uint8_t *data, size_t length);

    /**
     * @brief 从升级包中提取文件
     * @param data 数据指针
     * @param length 数据长度
     * @return 是否成功
     */
    bool ExtractFilesFromUpgradePackage(const uint8_t *data, size_t length);

    /**
     * @brief 解析CPIO头部
     * @param header 头部数据
     * @param fileInfo 输出的文件信息
     * @return 是否成功
     */
    bool ParseCpioHeader(const uint8_t *header, CpioFileInfo &fileInfo);

    /**
     * @brief 解析CPIO头部中的十六进制字段
     * @param hexStr 十六进制字符串
     * @param length 字符串长度
     * @return 转换后的数字
     */
    uint32_t ParseCpioHexField(const char *hexStr, size_t length);

    /**
     * @brief 处理当前文件的数据内容
     * @param data 数据指针
     * @param availableBytes 可用数据字节数
     * @return 处理的字节数
     */
    size_t ProcessCurrentFileContent(const uint8_t *data, size_t availableBytes);

    /**
     * @brief 解析描述文件（可选）
     * @param content 文件内容
     * @return 是否成功
     */
    bool ParseSwDescription(const std::string &content);

    /**
     * @brief 从描述信息中提取分区配置（可选）
     * @return 是否成功
     */
    bool ConfigurePartitionsFromSwDescription();


    /**
     * @brief 流式写入镜像数据
     * @param imageType 镜像类型 (kernel/rootfs/assets)
     * @param data 数据指针
     * @param length 数据长度
     * @param isLastChunk 是否最后一块数据
     * @return 是否成功
     */
    bool WriteImageStream(const std::string &imageType, const uint8_t *data, size_t length, bool isLastChunk);

    /**
     * @brief 确定当前活动分区和目标分区
     * @return 是否成功
     */
    bool DetermineTargetPartitions();

    /**
     * @brief 更新进度
     * @param percentage 进度百分比
     * @param status 状态描述
     */
    void UpdateProgress(int percentage, const std::string &status);
    
    /**
     * @brief 根据实际文件大小调整进度计算的总大小
     */
    void AdjustProgressTotalSize();

    /**
     * @brief 延迟重启系统
     * @param delaySeconds 延迟秒数
     */
    void ScheduleDelayedReboot(int delaySeconds = 10);

    /**
     * @brief 后台重启任务线程函数
     * @param delaySeconds 延迟秒数
     */
    void DelayedRebootTask(int delaySeconds);

private:
    std::unique_ptr<UBootEnvManager> ubootEnvManager_;
    std::unique_ptr<MtdManager> mtdManager_;
    std::unique_ptr<UbiManager> ubiManager_;

    std::mutex mutex_;
    std::atomic<bool> isRunning_;
    std::atomic<bool> upgradeInProgress_;

    // 升级包解析状态
    enum class UpgradePhase {
        WAITING_HEADER,      // 等待CPIO头部
        PARSING_DESCRIPTION, // 解析描述文件
        WRITING_IMAGES       // 写入镜像数据
    };
    UpgradePhase upgradePhase_;

    // CPIO文件解析状态
    enum class CpioParsingState {
        WAITING_HEADER,      // 等待CPIO头部
        WAITING_FILENAME,    // 等待文件名
        PROCESSING_FILE_DATA // 处理文件数据
    };
    CpioParsingState cpioParsingState_;
    size_t cpioStreamPosition_; // 当前处理到的流位置

    std::string currentStatus_;
    size_t streamBufferSize_; // 动态缓冲区大小

    // 流式处理状态
    std::vector<uint8_t> streamBuffer_;                                   // 流式数据缓冲区
    std::map<std::string, uint64_t> partitionWritePositions_;             // 各分区物理写入位置
    std::map<std::string, uint64_t> partitionLogicalWritePositions_;      // 各分区逻辑写入位置（排除坏块）
    std::map<std::string, bool> partitionWriteInitialized_;               // 分区写入是否已初始化
    std::map<std::string, std::vector<uint8_t>> partitionAlignmentBuffers_; // 各分区对齐缓冲区

    CpioFileInfo currentFile_;

    std::string swDescription_;
    nlohmann::json upgradeConfig_;
    bool descriptionParsed_ = false;                      // 是否已解析描述文件
    std::map<std::string, std::string> imageFilenameToType_; // 文件名 -> 镜像类型(kernel/rootfs/assets)

    // 分区信息
    std::string currentSlot_; // 当前活动分区 (A/B)
    std::string targetSlot_;  // 目标分区 (A/B)

    // 升级进度相关
    std::map<std::string, uint64_t> expectedFileSizes_;  // 预期文件大小（来自描述文件或设备信息）
    std::map<std::string, uint64_t> actualFileSizes_;    // 实际文件大小（从CPIO头部获取）
    uint64_t totalExpectedSize_;  // 总预期大小（初始基于分区大小）
    uint64_t totalActualSize_;    // 总实际大小（基于真实文件大小）
    uint64_t totalWrittenSize_;   // 总已写入大小
    bool progressSizeAdjusted_;   // 是否已调整进度大小

    // 进度回调
    UpgradeProgressCallback progressCallback_;

    // 状态和错误信息
    std::string lastError_;

    // 进度日志节流
    int lastLoggedProgressBucket_ = -1; // 上一次打印的10%档位（-1表示未打印过）
    bool firstProgressLogged_ = false;  // 是否已打印首次进度

    // 配置常量
    static constexpr size_t DEFAULT_STREAM_BUFFER_SIZE = 2 * 1024 * 1024; // 默认2MB流缓冲区
    static constexpr size_t MIN_STREAM_BUFFER_SIZE = 1 * 1024 * 1024;     // 最小1MB流缓冲区
    static constexpr size_t MAX_STREAM_BUFFER_SIZE = 4 * 1024 * 1024;     // 最大4MB流缓冲区
    static constexpr size_t MIN_ERASE_SIZE = 4 * 1024;                    // 最小擦除大小4KB
};

} // namespace Upgrade
} // namespace El
